#pragma once

#include <iostream>
#include <vector>
#include <algorithm>
#include <cmath>
#include <queue>
#include <unordered_set>
#include <memory>
#include <stack>
#include <execution>
#include <tbb/concurrent_vector.h>
#include <tbb/parallel_for.h>
#include <tbb/blocked_range.h>


#include "mltree/entities/mlnode.hpp"

namespace MetricSpaceBenchmark::MLIndex::MLTree {

    template<typename T>
    class BuildLayer {
    public:
        std::vector <Point<T>> centers;
        std::vector <std::unique_ptr<MLTreeNode<T>>> nodes;
        bool isLeafLayer{};

        BuildLayer() = default;

        BuildLayer(std::vector <Point<T>> centers, std::vector <std::unique_ptr<MLTreeNode<T>>> nodes, bool isLeaf)
                : centers(std::move(centers)), nodes(std::move(nodes)), isLeafLayer(isLeaf) {}

        BuildLayer(BuildLayer &&) noexcept = default;
    };

    template<typename T>
    class MLTree {
    private:
        std::unique_ptr <MLTreeNode<T>> root;

    public:
        MLTree() : root(nullptr) {}

        ~MLTree() = default;

        MLTree(MLTree &&) noexcept = default;

        MLTree &operator=(MLTree &&) noexcept = default;

        void printTree(size_t &count) const {
            if (root) {
                root->printNode(count);
            }
        }

        // 步骤 1：构建候选节点集 E
        void buildCandidateNodes(BuildLayer<T> &currentLayer, unsigned long M, bool isLeafLayer,
                                 std::vector <std::unique_ptr<MLTreeNode<T>>> &E);

        // 步骤 2：选择子集 A
        void selectSubsetA(const std::vector <std::unique_ptr<MLTreeNode<T>>> &E,
                           std::vector <std::unique_ptr<MLTreeNode<T>>> &A);

        // 步骤 3：非递归构建 M-Tree
        void buildMLTree(const std::vector <Point<T>> &objects, unsigned long M);

        // 范围查询函数
        void rangeQuery(const Point<T> &queryPoint, DIS_TYPE queryRadius, std::vector <Point<T>> &result);
    };

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    // 步骤 1：构建候选节点集 E
    template<typename T>
    void MLTree<T>::buildCandidateNodes(BuildLayer<T> &currentLayer, unsigned long M, bool isLeafLayer,
                                        std::vector <std::unique_ptr<MLTreeNode<T>>> &E) {

        const auto &centers = currentLayer.centers; // 从当前层获取中心点集合
        E.clear();
        std::vector<bool> used(centers.size(), false);

        // 生成点对距离优先队列（优化为并行计算）
        tbb::concurrent_vector <PointPair> pairs;
        tbb::parallel_for(tbb::blocked_range<size_t>(0, centers.size()),
                          [&](const tbb::blocked_range <size_t> &r) {
                              for (size_t i = r.begin(); i < r.end(); ++i) {
                                  for (size_t j = i + 1; j < centers.size(); ++j) {
                                      pairs.push_back({i, j, centers[i].distance(centers[j])});
                                  }
                              }
                          });

        std::priority_queue <PointPair, std::vector<PointPair>, std::greater<>> pq(std::greater<>(),
                                                                                   std::vector<PointPair>(pairs.begin(),
                                                                                                          pairs.end()));
        // 合并节点逻辑
        while (!pq.empty()) {
            const auto [i, j, dist] = pq.top();
            pq.pop();

            if (used[i] && used[j]) continue;

            bool merged = false, isFlag = true;
            for (auto &node: E) {
                const auto &nodePoints = node->getPoints();
//                const bool has_i = nodePoints.contains(centers[i]);
//                const bool has_j = nodePoints.contains(centers[j]);
                const bool has_i = nodePoints.count(centers[i]) > 0;
                const bool has_j = nodePoints.count(centers[j]) > 0;

                if (has_i || has_j) {
                    const auto &target = has_i ? centers[j] : centers[i];
                    if (node->addEntryIfPossible(target, M)) {
                        used[i] = used[j] = true;
                        merged = true;
                        break;
                    } else {
                        isFlag = false;
                    }
                }
            }

            if (!merged && isFlag) {
                if (isLeafLayer) {
                    std::vector <DBEntry<T>> newEntries;
                    newEntries.emplace_back(centers[i].oId, centers[i].keyObject);
                    newEntries.emplace_back(centers[j].oId, centers[j].keyObject);
                    E.emplace_back(std::make_unique<MTreeLeaf<T>>(newEntries));
                } else {
                    std::vector <Point<T>> newPoints{centers[i], centers[j]};
                    // 创建内部节点（需要携带子树信息）
                    std::vector <RoutingObject<T>> routingEntries;
                    routingEntries.reserve(newPoints.size());
                    for (const auto &p: newPoints) {
                        // 查找对应的子树（来自上一层的nodes）
                        auto it = std::find_if(currentLayer.nodes.begin(), currentLayer.nodes.end(),
                                               [&p](const auto &node) {
                                                   return node->getCenter().oId == p.oId; // 通过唯一ID匹配
                                               });

                        if (it != currentLayer.nodes.end()) {
                            routingEntries.emplace_back(p.oId, p.keyObject, DBL_MAX, p.dis, (*it)->clone());
                        }
                    }
                    E.emplace_back(std::make_unique<MTreeInternal<T>>(std::move(routingEntries)));
                }
                used[i] = used[j] = true;
            }
        }

        // 处理未合并的孤立点
//        for (size_t i = 0; i < centers.size(); ++i) {
//            if (!used[i]) {
//                if (isLeafLayer) {
//                    E.emplace_back(std::make_unique<MTreeLeaf<T>>(
//                            std::vector<DBEntry<T>>{DBEntry<T>(centers[i].oId, centers[i].keyObject)}));
//                } else {
//                    // 查找对应的子树
//                    auto it = std::find_if(currentLayer.nodes.begin(), currentLayer.nodes.end(),
//                                           [&](const auto &node) { return node->getCenter().oId == centers[i].oId; });
//
//                    if (it != currentLayer.nodes.end()) {
//                        E.emplace_back(
//                                std::make_unique<MTreeInternal<T>>(std::vector<RoutingObject<T>>{RoutingObject<T>(
//                                        centers[i].oId, centers[i].keyObject, DBL_MAX, centers[i].dis,
//                                        (*it)->clone())}));
//                    }
//                }
//            }
//        }

        // 并行更新所有节点参数
        std::for_each(std::execution::par, E.begin(), E.end(), [](auto &node) { node->updateCenterAndRadius(); });
    }

    // 步骤 2：选择子集 A
    template<typename T>
    void MLTree<T>::selectSubsetA(const std::vector <std::unique_ptr<MLTreeNode<T>>> &E,
                                  std::vector <std::unique_ptr<MLTreeNode<T>>> &A) {

        A.clear();
        if (E.empty()) return;

        // 按覆盖效率排序（面积/新增覆盖点数）
        std::vector < std::pair < double, std::unique_ptr < MLTreeNode<T>>>> sortedNodes;
        sortedNodes.reserve(E.size());

        std::unordered_set <Point<T>, PointHash<T>, PointEqual<T>> coveredPoints;

        // 预计算所有节点的覆盖效率
        for (const auto &node: E) {
            const auto points = node->getPoints();
            size_t newCoverage = std::count_if(points.begin(), points.end(),
                                               [&](const auto &p) { return !coveredPoints.count(p); });

            if (newCoverage > 0) {
                double efficiency = node->area() / newCoverage;
                sortedNodes.emplace_back(efficiency, node->clone());
            }
        }

        // 按效率升序排序
        std::sort(sortedNodes.begin(), sortedNodes.end(),
                  [](const auto &a, const auto &b) { return a.first < b.first; });

        // 贪心选择
        coveredPoints.clear();
        for (auto &[eff, node]: sortedNodes) {
            const auto points = node->getPoints();
            size_t newCount = std::count_if(points.begin(), points.end(),
                                            [&](const auto &p) { return !coveredPoints.count(p); });

            if (newCount > 0) {
                coveredPoints.insert(points.begin(), points.end());
                A.emplace_back(std::move(node));
            }
        }
    }

    // 步骤 3：非递归构建 M-Tree
    template<typename T>
    void MLTree<T>::buildMLTree(const std::vector <Point<T>> &objects, unsigned long M) {

        std::stack <BuildLayer<T>> layers;

        // 初始层：叶子节点层
        layers.push(BuildLayer<T>{
                objects,
                std::vector < std::unique_ptr < MLTreeNode<T>>>{}, // 显式构造空vector
                true
        });

        while (!layers.empty()) {
            auto currentLayer = std::move(layers.top());
            layers.pop();

            // 步骤1：构建候选节点集E
            std::vector <std::unique_ptr<MLTreeNode<T>>> E;
            buildCandidateNodes(currentLayer, M, currentLayer.isLeafLayer, E);

            // 步骤2：选择子集A
            std::vector <std::unique_ptr<MLTreeNode<T>>> A;
            selectSubsetA(E, A);

            if (A.size() == 1) {
                root = std::move(A.front());
                continue;
            }

            // 准备下一层数据
            std::vector <Point<T>> nextLayer;
            nextLayer.reserve(A.size());

            // 创建父节点
            std::vector <RoutingObject<T>> routingEntries;
            routingEntries.reserve(A.size());

            for (auto &node: A) {
                auto center = node->getCenter();
                auto radius = node->getRadius();

                // 创建路由条目：中心点 + 覆盖半径 + 子树
                routingEntries.emplace_back(center.oId, center.keyObject, DBL_MAX, radius, std::move(node));

                // 暂时将 center 的 dis 视为覆盖半径
                center.dis = radius;

                nextLayer.emplace_back(center);
            }

            // 创建父节点
            auto parentNode = std::make_unique<MTreeInternal<T>>(std::move(routingEntries));
            parentNode->updateCenterAndRadius();

            // 压入新层
            std::vector <std::unique_ptr<MLTreeNode<T>>> nextNodes;
            nextNodes.emplace_back(std::move(parentNode)); // 显式移动

            layers.push(BuildLayer<T>{std::move(nextLayer), std::move(nextNodes), false});
        }
    }

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
    // 范围查询函数
    template<typename T>
    void MLTree<T>::rangeQuery(const Point<T> &queryPoint, DIS_TYPE queryRadius, std::vector <Point<T>> &result) {
        if (root) {
            dynamic_cast<MTreeInternal<T> *>(root.get())->getEntries().front().subtree->rangeQuery(queryPoint,
                                                                                                   queryRadius,
                                                                                                   result);
        }
    }
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
}